'''
@author: zhangkai
@license: (C) Copyright 2017-2023
@contact: jeffcobile@gmail.com
@Software : PyCharm
@file: darknet.py
@time: 2020-06-17 10:04:03
@desc: 
'''
import torch
from jjzhk.config import ZKCFG
from ELib.backbone.backbone_zoo import BACKBONE_ZOO, BackboneSeg


@BACKBONE_ZOO.register()
def darknet19(cfg:ZKCFG):
    return Darknet(cfg, cfg.MODEL.BACKBONE)


@BACKBONE_ZOO.register()
def darknet53(cfg:ZKCFG):
    return Darknet(cfg, cfg.MODEL.BACKBONE)


class Darknet(BackboneSeg):
    def __init__(self, cfg, type):
        super(Darknet, self).__init__(cfg)
        self.cfg = cfg
        self.type = type

    def _create_network_(self):
        if self.type.endswith('19'):
            return self._create_19_()
        elif self.type.endswith('53'):
            return self._create_53_()

    def _create_19_(self):
        layer = []
        layer += [Darknet_conv_bn(3, 32, 1)]
        layer += [torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)]
        layer += [Darknet_conv_bn(32, 64, 1)]
        layer += [torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)]

        layer += [Darknet_conv_block(64, 128, 1)]
        layer += [Darknet_conv_block(128, 128, 1)]

        layer += [torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)]

        layer += [Darknet_conv_block(128, 256, 1)]
        layer += [Darknet_conv_block(256, 256, 1)]
        layer += [torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)]

        layer += [Darknet_conv_block(256, 512, 1)]
        layer += [Darknet_conv_block(512, 512, 1)]
        layer += [Darknet_conv_block(512, 512, 1)]
        layer += [torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)]

        layer += [Darknet_conv_block(512, 1024, 1)]
        layer += [Darknet_conv_block(1024, 1024, 1)]
        layer += [Darknet_conv_block(1024, 1024, 1)]

        return layer

    def _create_53_(self):
        layer = []
        layer += [Darknet_conv_bn(3, 32, 1)]
        layer += [Darknet_conv_block(32, 64, stride=1, darknettype=53)]
        layer += [Darknet_conv_block(64, 64, stride=1, darknettype=53)]
        layer += [Darknet_conv_block(64, 128, stride=1, darknettype=53)]
        layer += [Darknet_conv_block(128, 128, stride=1, darknettype=53)]
        layer += [Darknet_conv_block(128, 128, stride=1, darknettype=53)]

        layer += [Darknet_conv_block(128, 256, stride=1, darknettype=53)]
        layer += [Darknet_conv_block(256, 256, stride=1, darknettype=53)]
        layer += [Darknet_conv_block(256, 256, stride=1, darknettype=53)]

        layer += [Darknet_conv_block(256, 256, stride=1, darknettype=53)]
        layer += [Darknet_conv_block(256, 256, stride=1, darknettype=53)]
        layer += [Darknet_conv_block(256, 256, stride=1, darknettype=53)]
        layer += [Darknet_conv_block(256, 256, stride=1, darknettype=53)]

        layer += [Darknet_conv_block(256, 256, stride=1, darknettype=53)]
        layer += [Darknet_conv_block(256, 256, stride=1, darknettype=53)]
        layer += [Darknet_conv_block(256, 512, stride=1, darknettype=53)]

        layer += [Darknet_conv_block(512, 512, stride=1, darknettype=53)]
        layer += [Darknet_conv_block(512, 512, stride=1, darknettype=53)]
        layer += [Darknet_conv_block(512, 512, stride=1, darknettype=53)]
        layer += [Darknet_conv_block(512, 512, stride=1, darknettype=53)]
        layer += [Darknet_conv_block(512, 512, stride=1, darknettype=53)]

        layer += [Darknet_conv_block(512, 512, stride=1, darknettype=53)]
        layer += [Darknet_conv_block(512, 512, stride=1, darknettype=53)]
        layer += [Darknet_conv_block(512, 512, stride=1, darknettype=53)]

        layer += [Darknet_conv_block(512, 1024, stride=1, darknettype=53)]

        layer += [Darknet_conv_block(1024, 1024, stride=1, darknettype=53)]
        layer += [Darknet_conv_block(1024, 1024, stride=1, darknettype=53)]
        layer += [Darknet_conv_block(1024, 1024, stride=1, darknettype=53)]
        layer += [Darknet_conv_block(1024, 1024, stride=1, darknettype=53)]
        return layer


class Darknet_conv_bn(torch.nn.Module):
    def __init__(self, inp, oup, stride):
        super(Darknet_conv_bn, self).__init__()
        self.conv = torch.nn.Sequential(
            torch.nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
            torch.nn.BatchNorm2d(oup),
            torch.nn.LeakyReLU(0.1, inplace=True),
        )
        self.depth = oup

    def forward(self, x):
        return self.conv(x)


class Darknet_conv_block(torch.nn.Module):
    def __init__(self, inp, oup, stride, expand_ratio=0.5, darknettype=19):
        super(Darknet_conv_block, self).__init__()
        self.darknet_type = darknettype
        self.use_res_connect = stride == 1 and inp == oup
        if self.use_res_connect:
            depth = int(oup*expand_ratio)
            self.conv = torch.nn.Sequential(
                torch.nn.Conv2d(inp, depth, 1, 1, bias=False),
                torch.nn.BatchNorm2d(depth),
                torch.nn.LeakyReLU(0.1, inplace=True),
                torch.nn.Conv2d(depth, oup, 3, stride, 1, bias=False),
                torch.nn.BatchNorm2d(oup),
                torch.nn.LeakyReLU(0.1, inplace=True),
            )
        else:
            self.conv = torch.nn.Sequential(
                torch.nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
                torch.nn.BatchNorm2d(oup),
                torch.nn.LeakyReLU(0.1, inplace=True),
            )
        self.depth = oup

    def forward(self, x):
        if self.darknet_type == 19:
            return self.conv(x)
        else:
            if self.use_res_connect:
                return x + self.conv(x)
            else:
                return self.conv(x)